import os

import numpy
import torch
import torch.nn as nn
from torchvision import models
from torchvision.transforms import transforms
import h5py
import cv2
import sys
import numpy as np
import matplotlib as plt
from PIL import Image

def img_process(h5file):
    f = h5py.File(h5file, 'r')  # 打开h5文件
    class_tiff = f['class_tiff'][:]
    defect_tiff = f['defect_tiff'][:]
    f.close()

    transform_test = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.40272543], std=[0.13901867])
    ])

    image_class = Image.fromarray(class_tiff).convert("L")
    image_class.save(os.path.join(output_dir, "origin_class.jpg"))
    image_class = transform_test(image_class)
    image_class = torch.unsqueeze(image_class, dim=0)

    image_defect = Image.fromarray(defect_tiff).convert("L")
    image_defect.save(os.path.join(output_dir, "origin_defect.jpg"))
    image_defect = transform_test(image_defect)
    image_defect = torch.unsqueeze(image_defect, dim=0)

    return image_class, image_defect


def backward_hook(module, grad_in, grad_out):
    grad_block.append(grad_out[0].clone())

def forward_hook(module, input, output):
    fmap_block.append(output)


def cam_show_img(img, feature_map, grads, out_dir):
    img = torch.from_numpy(img)
    img = img.permute(1, 2, 0)
    img = np.array(img)
    H, W, C = img.shape
    cam = np.zeros(feature_map.shape[1:], dtype=np.float32)
    grads = grads.reshape([grads.shape[0], -1])
    weights = np.mean(grads, axis=1)
    for i, w in enumerate(weights):
        cam += w * feature_map[i, :, :]

    # print(cam.shape)#(4, 4)
    cam = np.maximum(cam, 0)
    cam = cam / cam.max()
    cam = cv2.resize(cam, (H, W))

    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)


    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
    cam_img = 0.3 * heatmap + 0.7 * img
    path_cam_img = os.path.join(out_dir, "cam.jpg")
    cv2.imwrite(path_cam_img, cam_img)

if __name__ == '__main__':
    categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch',
                  'Pattern_fail',
                  'PR_peeling', 'Seam', 'Reference', 'Surface_particle', 'Burried_particle', 'Cu_diffuse',
                  'Prelayer_defect_developed',
                  'Void', 'Residue', 'Scratch']
    categories_to_id = dict((c, i) for i, c in enumerate(categories))
    id_to_categories = dict((v, k) for k, v in categories_to_id.items())

    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    grad_block = list()
    fmap_block = list()
    output_dir = r"/home/yeadc/Documents/cxj/Classify-Seg/task_classify/task/heatmap/temp"
    h5file = r"/home/yeadc/Documents/cxj/Classify-Seg/task_classify/task/heatmap/h5/Residue_FY-A100_BE_ASI_lFFJ310940_w13_338_Topography3.h5"
    img_class, img_defect = img_process(h5file)
    img_class = img_class.to(DEVICE)
    img_defect = img_defect.to(DEVICE)


    model_weight = r"/home/yeadc/Documents/cxj/Classify-Seg/task_classify/task/checkpoint/swinv2_base_patch8_240426/model_92_87.32.pth"
    model = torch.load(model_weight)
    model.eval()
    print(model)


    model.layers[3].register_forward_hook(forward_hook)
    model.layers[3].register_backward_hook(backward_hook)



    #forward
    output = model(img_class)

    #backward
    model.zero_grad()
    loss_fn = nn.CrossEntropyLoss()
    target = torch.unsqueeze(torch.from_numpy(np.array(categories_to_id['Residue'])), dim=0)
    target = target.to(DEVICE)
    loss = loss_fn(output, target)
    loss.backward()

    #generate cam
    grads_val = grad_block[0].cpu().data.numpy().squeeze()
    fmap = fmap_block[0].cpu().data.numpy().squeeze()

    f = h5py.File(h5file, 'r')  # 打开h5文件
    img_class = np.array(f['class_tiff'][:])
    img_defect = np.array(f['defect_tiff'][:])
    f.close()
    img_class = img_class[np.newaxis, :]
    img_defect = img_defect[np.newaxis, :]
    cam_show_img(img_class, fmap, grads_val, output_dir)



